#coding:utf8

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from matplotlib.cbook import flatten
import math

qpsk = [1+1j, 1-1j, -1+1j, -1-1j]
psk8 = [(0,np.sqrt(2)), (np.sqrt(2),0), (-np.sqrt(2),0), (0,-np.sqrt(2)), (1,1),(1,-1),(-1,1),(-1,-1)]
psk8_complex = map(lambda p: p[0] + p[1]*1j, psk8)
pi_8 = np.sin(math.pi/8)
pi3_8 = np.sin(math.pi*3/8)
pi2_8 = np.sin(math.pi/4)
psk16 = [(1, 0), (pi3_8, pi_8), (pi2_8, pi2_8), (pi_8, pi3_8),
        (0, 1), (-pi_8, pi3_8), (-pi2_8,pi2_8), (-pi3_8,pi_8),
        (-1, 0), (-pi3_8, -pi_8), (-pi2_8, -pi2_8), (-pi_8, -pi3_8),
        (0, -1), (pi_8, -pi3_8), (pi2_8, -pi2_8), (pi3_8, -pi_8)
        ]
psk16 = map(lambda p: (p[0]*np.sqrt(2), p[1]*np.sqrt(2)), psk16)
psk16_complex = map(lambda p: p[0] + p[1]*1j, psk16)


def basis(l,i):
    res = [0 for _ in range(l)]
    res[i-1] = 1
    return res

#QPSK信号
def generateData(M,T,dB,L,qn=4,flag=None):
    def yinshe(x):
        if qn == 4:
            return basis(4, qpsk.index(x))
        if qn == 8:
            return basis(8, psk8_complex.index(x))
        if qn == 16:
            return basis(16, psk16_complex.index(x))
    #QPSK信号
    if qn == 4:
        TxS = np.sign(np.random.rand(M) * 2 - 1) + 1j*np.sign(np.random.rand(M) * 2 - 1) #30000
    if qn == 8:
        tmp = map(int, np.random.rand(M) * 8)
        TxS = map(lambda x: psk8[x][0] + 1j * psk8[x][1],tmp)
    if qn == 16:
        tmp = map(int, np.random.rand(M) * 16)
        TxS = map(lambda x: psk16[x][0] + 1j * psk16[x][1],tmp)
    ch = [0.0410+0.0109j,0.0495+0.0123j,0.0672+0.0170j,0.0919+0.0235j,
     0.7920+0.1281j,0.3960+0.0871j,0.2715+0.0498j,0.2291+0.0414j,0.1287+0.0154j,
     0.1032+0.0119j]
    ch = ch / np.linalg.norm(ch)
    x = signal.fftconvolve(ch,TxS)[:M]   #信道卷积 x.shape = (30000,)
    #noise
    n=np.random.randn(1,M)+1j*np.random.randn(1,M);
    n=n/np.linalg.norm(n)*pow(10,(-dB/20.0))*np.linalg.norm(x);
    x = x + n
    x = x.ravel()
    K = M-L-1 #29987
    X = []
    for i in range(K):
        X.append(x[i+L+1:i:-1])
    X = np.array(X).T  # (13,29987)
    TxS = TxS[L:M-6]
    Y = X[:,5:]

    if flag == 'rnn':
        Tx = [yinshe(a) for a in np.array(TxS).T]
        Tx = np.array(Tx)
        Y = np.array([list(flatten(zip(np.real(y), np.imag(y)))) for y in Y.T])
        return Y,Tx,x
    if flag == 'cnn':
        Tx = [yinshe(a) for a in np.array(TxS).T]
        Tx = np.array(Tx)
        Y = np.array([list(zip(np.real(y), np.imag(y))) for y in Y.T])
        return Y,np.array(Tx).T,x
    #Y为训练集（接收机收到的数据），TxS为target(发送序列), x为接收端接收到的数据(用来画图比较)
    return Y,np.array(TxS).T,x

def calc_distance(p1, p2):
    return np.sqrt(pow(p1[0]-p2[0],2) + pow(p1[1]-p2[1],2))


if __name__ == '__main__':
    X,Y,_ = generateData(3000,2000,10,12, qn=16,flag='rnn')
    print Y
    X,Y,_ = generateData(3000,2000,10,12, flag='cnn')
    print X.shape, Y.shape

    # plt.scatter(x.real, x.imag)
    # plt.show()
